import io
import os
import glob
import fitz
from time import perf_counter
import PIL.Image

DEV_MODE = True  # 开发或测试模式，调试完，改成False


def file_name_to_folder(file, src_folder, des_folder):
    # 如果扩展名有其他比如大写情况，后面可再加.replace('.PDF', '')
    return file.replace('.pdf', '').replace(src_folder, des_folder)


def pdf_to_images(file, file_renamer, jpg2000):
    folder = file_renamer(file)
    os.makedirs(folder, exist_ok=True)
    doc = fitz.open(file)
    extension = 'jp2' if jpg2000 else 'jpg'
    image_format_str = 'jpeg2000' if jpg2000 else 'jpeg'
    for page in doc:
        if DEV_MODE:
            if page.number > 10:
                return
        image_file = os.path.join(folder, f"{page.number + 1}.{extension}")
        # 默认跳过已经生成过的图片
        if not DEV_MODE:
            if os.path.isfile(image_file):
                continue
        try:
            pix = page.get_pixmap(dpi=600)  # dpi默认值96，越小越快
            img = PIL.Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
            img = img.reduce(2)  # 缩小倍数
            img = img.convert('L')  # 单色模式，图片更小，1表示每个像素纯黑白，L表示255个灰度
            img.save(image_file,
                     optimize=True,  # 最高压缩率
                     quality=75,  # 范围1-95，默认值75
                     dpi=(300, 300),  # 越低尺寸越小
                     format=image_format_str
                     )
            print(image_file)
        except Exception as e:
            print(f'{image_file} {e}')


def convert_pdfs_in_folder(pdf_folder, file_renamer, jpg2000):
    for file in glob.glob(rf'{pdf_folder}\**\*.pdf', recursive=True):
        pdf_to_images(file, file_renamer, jpg2000)


if __name__ == '__main__':
    # 这后就两个参数是pdf的某个路径替换为图片路径，可以灵活指定替换规则
    t_start = perf_counter()
    file_renamer = lambda file: file_name_to_folder(file, r'D:\dl\testpdfs',
                                                    r'D:\dl\jpgs')
    # 默认为普通jpg，设置为True则生成jpg2000格式扩展名jp2
    jpg2000 = False
    # jpg2000 = True
    convert_pdfs_in_folder(r'D:\dl\testpdfs', file_renamer, jpg2000
                           )
    t_end = perf_counter()
    # 用了多少秒，测试用
    if DEV_MODE:
        print(f'用了{t_end - t_start}秒')
